"""
NBA球队赛季场均得分分析应用（并行计算优化版本）

功能描述：
    使用PySpark对NBA比赛数据进行并行批处理分析，计算每支球队在各赛季的场均得分。
    通过分区策略和并行计算提高处理效率。

优化策略：
    1. 按赛季ID和球队进行分区，实现数据的并行处理
    2. 使用缓存机制减少重复计算
    3. 实现分布式计算，提高处理效率
"""

from pyspark.sql import SparkSession
from pyspark.sql.functions import sum, count, round, col, concat_ws
import os

def calculate_team_season_avg_points():
    """计算NBA每支球队的赛季场均得分（并行优化版本）
    优化要点：
    1. 设置合适的分区数
    2. 使用缓存减少重复计算
    3. 并行处理各球队数据
    4. 合理使用shuffle操作
    """
    
    # 第一步：配置Spark环境(设置分区数和并行数为4)
    spark = SparkSession.builder \
        .appName("NBA Team Season Average Points Parallel") \
        .config("spark.sql.shuffle.partitions", 4) \
        .config("spark.default.parallelism", 4) \
        .getOrCreate()

    # 设置日志级别为ERROR，减少不必要的日志输出
    spark.sparkContext.setLogLevel("ERROR")

    try:
        # 第二步：数据加载和分区优化
        csv_path = "../data/game_processed.csv"
        
        # 读取CSV文件并进行分区优化
        df = spark.read \
            .option("header", "true") \
            .option("inferSchema", "true") \
            .csv(csv_path) \
            .repartition(4, "season_id")  # 按season_id列重新分区，提高并行度
        
        # 缓存数据，避免重复计算
        df.cache()

        # 第三步：并行处理主场数据
        home_stats = df.groupBy("season_id", "team_abbreviation_home").agg(
            sum("pts_home").alias("total_points"),
            count("*").alias("num_games")
        ).select(
            col("season_id"),
            col("team_abbreviation_home").alias("team"),
            col("total_points"),
            col("num_games")
        )

        # 第四步：并行处理客场数据
        away_stats = df.groupBy("season_id", "team_abbreviation_away").agg(
            sum("pts_away").alias("total_points"),
            count("*").alias("num_games")
        ).select(
            col("season_id"),
            col("team_abbreviation_away").alias("team"),
            col("total_points"),
            col("num_games")
        )

        # 第五步：合并主客场数据
        all_stats = home_stats.union(away_stats)

        # 第六步：计算每支球队的场均得分
        team_stats = all_stats.groupBy("season_id", "team").agg(
            sum("total_points").alias("season_total_points"),
            sum("num_games").alias("season_games")
        )

        # 第七步：计算最终结果
        result = team_stats.select(
            "season_id",
            "team",
            round(
                col("season_total_points") / col("season_games"),
                2
            ).alias("avg_points_per_game")
        ).orderBy("season_id", "team")

        # 第八步：结果输出优化
        output_path = "../data/team_season_avg_points_parallel.csv"
        result.coalesce(1).write \
            .mode("overwrite") \
            .option("header", "true") \
            .csv(output_path)

        # 第九步：显示结果和性能统计
        print("\n每支球队的赛季场均得分：")
        result.show(truncate=False)

        # 显示执行计划，用于性能分析
        print("\n执行计划：")
        result.explain()

        # 第十步：计算和显示球队得分分布
        season_summary = result.groupBy("season_id").agg(
            round(max("avg_points_per_game"), 2).alias("max_points"),
            round(min("avg_points_per_game"), 2).alias("min_points"),
            round(sum("avg_points_per_game")/count("*"), 2).alias("avg_points")
        ).orderBy("season_id")

        print("\n赛季得分分布分析：")
        season_summary.show(truncate=False)

        # 第十一步：清理缓存
        df.unpersist()

    finally:
        # 第十二步：清理资源
        spark.stop()

if __name__ == "__main__":
    calculate_team_season_avg_points()
